/**============新的MatchLight(明暗交替)============
 * 作者：hqy
 * 创建时间：2020.2.9 00:27
    输入按照装甲板上点的顺序: 从左上角开始逆时针,
    |0        3|
    |1        2|   
 */
#include "aim_deps/LOG.hpp"
#include "aim_auto/ArmorPlate.hpp"
#include "aim_deps/Paths.hpp"

const float threshold = 0.15;

ArmorPlate::ArmorPlate(){
    for(int i=0; i<4 ;++i) points[i] = aim_deps::NULLPOINT2f;
    std::string path = model_path + "traced_model.pt";
    clf = torch::jit::load(path);
}

ArmorPlate::~ArmorPlate(){
    ;
}

void ArmorPlate::matchAll(
    const std::vector<cv::Point> &matches,
    std::vector<aim_deps::Light> &lights, 
    std::vector<aim_deps::Armor> &tar_list
)
{

    tar_list.clear();
    // #pragma omp parallel for num_threads(2)
    for(size_t i = 0 ; i < matches.size() ; ++i){
        const aim_deps::Light& l1 = lights[matches[i].x];
        const aim_deps::Light& l2 = lights[matches[i].y];
        std::vector<float> feature(5);
        feature[0] = l1.box.length;
        feature[1] = l2.box.length;
        cv::Point2f vec = l2.box.center - l1.box.center;
        float cosa = std::cos(-l1.box.angle * aim_deps::DEG2RAD), sina = std::sin(-l1.box.angle * aim_deps::DEG2RAD);
        feature[2] = cosa * vec.x - sina * vec.y;
        feature[3] = sina * vec.x + cosa * vec.y;
        feature[4] = l2.box.angle - l1.box.angle;
        at::Tensor ts = torch::tensor(feature).reshape({1, 5});
        at::Tensor res = clf.forward({ts}).toTensor().clamp(0, 1);
        float r1 = res.index({0, 0}).item().toFloat();
        float r2 = res.index({0, 1}).item().toFloat();
        int v1 = (r1 - r2) >= threshold;
        int v2 = (r2 - r1) >= threshold;
        if ((v1 | v2) && ((v1 & v2) == 0)) {        // v1 v2 中有1，说明识别到装甲板
            const aim_deps::Light& l1 = lights[matches[i].x];
            const aim_deps::Light& l2 = lights[matches[i].y];
            points[0] = l1.box.vex[0];
            points[1] = l1.box.vex[1];
            points[2] = l2.box.vex[1];
            points[3] = l2.box.vex[0];
            tar_list.emplace_back(points, 0, l1, l2);
            tar_list.back().valid = true;
            tar_list.back().Isbigarmor = (v2 == 1);
        }
    }
    filter(tar_list, lights);
}

void ArmorPlate::drawArmorPlates(cv::Mat &src, 
    const std::vector<aim_deps::Armor>& tar_list, const int optimal,const bool is_spin) const{
	char str[5];
    cv::line(src, cv::Point(aim_deps::VIDEO_WIDTH/2, 0), cv::Point(aim_deps::VIDEO_WIDTH/2, aim_deps::VIDEO_HEIGHT), cv::Scalar(255, 0, 0));
	cv::line(src, cv::Point(0, aim_deps::VIDEO_HEIGHT/2), cv::Point(aim_deps::VIDEO_WIDTH, aim_deps::VIDEO_HEIGHT/2), cv::Scalar(255, 0, 0));
    for (size_t i = 0; i< tar_list.size(); ++i) {
        if(tar_list[i].armor_number != -1){   //有意义的数字
            int line_width = tar_list[i].Isbigarmor ? 4 : 1;
            if(!tar_list[i].valid){       //无效装甲板使用黄色绘制
                for (int j = 0; j < 4; ++j){
                    // cv::line(src, tar_list[i].vertex[j], 
                    // tar_list[i].vertex[(j + 1) % 4], cv::Scalar(0, 255, 255), line_width);   
                }
            }                          //最佳装甲板使用红色绘制
            else if(i == size_t(optimal)) {
                for (int j = 0; j < 4; ++j){
                    cv::line(src, tar_list[i].vertex[j], 
                    tar_list[i].vertex[(j + 1) % 4], cv::Scalar(0, 0, 255), line_width);   
                }
                if(is_spin){
                cv::putText(src, "Spining", tar_list[i].vertex[3] + cv::Point2f(8, 4),
                    cv::FONT_HERSHEY_PLAIN, 1, cv::Scalar(255, 0, 255));
                }
            }
            else{                   //非最佳装甲板使用绿色绘制
                for (int j = 0; j < 4; ++j){
                    cv::line(src, tar_list[i].vertex[j],
                    tar_list[i].vertex[(j + 1) % 4], cv::Scalar(0, 255, 0), line_width);
                }
            }
            ///snprintf(str, 2, "%d", j);      //最佳装甲板位置x
	        ///cv::putText(src, str, tar_list[i].vertex[j]+cv::Point2f(2, 2),
	        ///    cv::FONT_HERSHEY_PLAIN, 1.1, cv::Scalar(0, 100, 255));
            snprintf(str, 5, "%d", tar_list[i].armor_number); //显示数字
            cv::putText(src, str, tar_list[i].vertex[2] + cv::Point2f(8, 4),
                cv::FONT_HERSHEY_PLAIN, 2, cv::Scalar(255, 255, 255));
            
        }
    }
}

int ArmorPlate::isIntersect(const aim_deps::Light& non_com1, const aim_deps::Light& non_com2, const aim_deps::Light& com) {
    cv::Point2f ncc1 = non_com1.box.center;
    cv::Point2f ncc2 = non_com2.box.center;
    cv::Point2f cc = com.box.center;
    cv::Point2f vec1 = ncc2 - ncc1;
    cv::Point2f vec2 = cc - ncc1;
    vec1 /= std::sqrt(vec1.dot(vec1));
    vec2 /= std::sqrt(vec2.dot(vec2));
    bool angle1_small = (vec1.dot(vec2) > -0.1736);            // 夹角小于100度

    cv::Point2f vec3 = cc - ncc2;
    vec1 *= -1;
    vec3 /= std::sqrt(vec3.dot(vec3));
    bool angle2_small = (vec1.dot(vec3) > -0.1736);            // 夹角小于100度
    if (angle1_small && angle2_small) return 0;
    if (angle1_small == false) 
        return 2; 
    if (angle2_small == false)
        return 1;
    return 0;
}   

struct ArmorCmpFunctor {
    bool operator()(const aim_deps::Armor* const a1, const aim_deps::Armor* const a2) const {
        return (std::abs(a1->left_light.box.angle - a1->right_light.box.angle) > 
            std::abs(a2->left_light.box.angle - a2->right_light.box.angle)
        );
    }
};

void ArmorPlate::filter(
    std::vector<aim_deps::Armor> &tar_list,
    const std::vector<aim_deps::Light>& lights
) const {
    size_t lights_size = lights.size();
    bool light_avl[lights_size];
    for(size_t i = 0; i<tar_list.size(); ++i){
        if(tar_list[i].valid == false) continue;
        for(size_t j = i+1; j< tar_list.size(); ++j){
            if(tar_list[j].valid == false) continue;
            std::pair<int, int> info;
            int judge = tar_list[i].collide(tar_list[j], info);
            if(judge == -1) continue;                  //装甲板无灯条冲突
            const aim_deps::Light& non_com1 = lights[info.first];
            const aim_deps::Light& non_com2 = lights[info.second];
            const aim_deps::Light& com = lights[judge];
            int intersect = isIntersect(non_com1, non_com2, com);
            if (intersect == 1) {
                tar_list[i].valid = false;
            }
            else if (intersect == 2) {
                tar_list[j].valid = false;
            }
        }
    }
    for (size_t i = 0; i < lights_size; i++) {
        light_avl[i] = true;
    }
    std::priority_queue<aim_deps::Armor*, std::vector<aim_deps::Armor*>, ArmorCmpFunctor> que;
    for (aim_deps::Armor& arm: tar_list) {
        if (arm.valid == true)
            que.emplace(&arm);
    }
    while (que.empty() == false) {
        aim_deps::Armor* ptr = que.top();
        int l_idx = ptr->left_light.index, r_idx = ptr->right_light.index;
        if (light_avl[l_idx] == false || light_avl[r_idx] == false) {
            ptr->valid = false;
        }
        else {
            light_avl[l_idx] = false;
            light_avl[r_idx] = false;
        }
        que.pop();
    }
}